'''
分析贡献度
'''
from numpy.lib.type_check import _imag_dispatcher
import torch
import torch.nn.functional as F
from torch.nn.modules import padding
from torch.serialization import check_module_version_greater_or_equal
import numpy as np
import os

def CreateFolder(folderpath):
    if not os.path.exists(folderpath):
        os.makedirs(folderpath)


def splitConvLayer(conv_modules):
    '''
    分离卷积层,构建不同的结果
    '''
    convlayers={}
    convlayers['conv1']={'conv':conv_modules['conv_1'],
                        'batchnorm':conv_modules['batchNorm2d_1'],
                        'ReLU':conv_modules['ReLU_1']}
    convlayers['conv2']={'conv':conv_modules['conv_2'],
                        'batchnorm':conv_modules['batchNorm2d_2'],
                        'ReLU':conv_modules['ReLU_2']}
    return convlayers

# --卷积运算结构模块------------------------------------------------------------
def Convolution(img_input,kernel_weight):
    '''
    针对单张影像具体的卷积运算
    args:
        img_imput:tensor 1,1,h,w
        kernel_weight:tensor 1,1,k_h,k_w  
    '''
    img_input_ex=torch.nn.ReplicationPad2d(1)(img_input) # 扩展
    result=F.conv2d(img_input_ex,kernel_weight,padding=0) # 确定运算的结果
    return result


def CalculateSingleConvoluation(img_input,kernel_weight):
    '''
    计算单个核权计算的结果
    args:
        img_input: tensor  1,c,h,w  用于运算的输入数据
        kernel_weight:tensor  1,c,h,w对应的权重核
    return：
        result： dict 此权重的单个波段与之对应的输出结果与最终的结果
        'contribution':计算贡献度的结果
                ''
    将img_input波段进行split，分别计算与之对应的波段计算结果
    '''
    b,c,h,w=img_input.shape
    kernel_b,kernel_c,kernel_h,kernel_w=kernel_weight.shape
    result={'kernel':kernel_weight.detach().cpu().numpy(),
            'kernel_result':{},
            'contribution':None,
            'output':None}
    result_sum=None
    for c_idx in range(c):
        img_input_c_idx=img_input[0,c_idx,:,:].reshape(1,1,h,w)
        kernel_weight_c_idx=kernel_weight[0,c_idx,:,:].reshape(1,1,kernel_h,kernel_w) #确定输出权重
        result_conv_c_idx=Convolution(img_input_c_idx,kernel_weight_c_idx)
        result['kernel_result'][c_idx]={'weight':kernel_weight_c_idx.detach().cpu().numpy(),
                                        'result':result_conv_c_idx.detach().cpu().numpy()}
        if result_sum is None:
            result_sum=result_conv_c_idx.detach().cpu()
        else:
            result_sum=result_sum+result_conv_c_idx.detach().cpu()
    # 贡献度计算
    convtribution_matrix=np.zeros((1,c,h,w))
    for c_idx in range(c):
        convtribution_matrix[0,c_idx,:,:]=result['kernel_result'][c_idx]['result']
    convtribution_matrix=convtribution_matrix/result_sum
    result['contribution']=convtribution_matrix
    # 结果验证
    #Compare_result=result_sum-result_conv.detach().cpu()
    #print(torch.sum(Compare_result))
    result['output']=result_sum.detach().cpu().numpy()
    return result


def CalulationConvolution(img_imput,kernel_weight):
    '''
    用于计算img_input，并记录每种权重核的运算结果
    args:
        img_input: tensor 1,c,h,w 输入用于运算的图片
        kernel_weight:tensor b,c,h,w 权重核
    return ：
        result:dict 记录返回值
                'img':记录当前用作运算的结果值
                'kernels'：记录当前权重情况
                'kernel_sig':记录单个权重计算的结果 
                
    '''
    b,c,h,w=img_imput.shape
    kernel_b,kernel_c,kernel_h,kernel_w=kernel_weight.shape
    result={'img':img_imput.detach().cpu().numpy(),
            'kernels':kernel_weight.detach().cpu().numpy(),
            'kernel_sig':{}}
    img_imput=img_imput.cuda().float()
    for b_idx in range(kernel_b):
        # 计算
        kernel_weight_b_idx=kernel_weight[b_idx,:,:,:].reshape(1,kernel_c,kernel_h,kernel_w)
        result['kernel_sig'][b_idx]=CalculateSingleConvoluation(img_imput,kernel_weight_b_idx)
    # 将单个波段整合成一个array
    result_np=np.ones((1,kernel_b,h,w))
    for b_idx in result['kernel_sig']:
        result_np[0,b_idx,:,:]=result['kernel_sig'][b_idx]['output']
    result['output']=result_np
    return result 

# -- 正则结构模块 -----
'''
注意正则模块有个别的参数需要在此时确定
'''

def BatchNorm(img_input,batchnorm_mould):
    '''
    正则运算
    '''
    img_input_out=batchnorm_mould(img_input)
    result={'img_input':img_input.detach().cpu().numpy(),
            'batchnorm_bias':batchnorm_mould.bias.detach().cpu().numpy(),
            'batchnorm_mean':batchnorm_mould.running_mean.detach().cpu().numpy(),
            'batchnorm_var':batchnorm_mould.running_var.detach().cpu().numpy(),
            'batchnorm_gamma':batchnorm_mould.weight.detach().cpu().numpy(),
            'output':img_input_out.detach().cpu().numpy()}
    return result

def ReLuCalucation(img_input,ReLU_mould):
    '''
    计算ReLu
    '''
    img_input_out=ReLU_mould(img_input)
    return img_input_out.detach().cpu().numpy()

# -- 计算单层结果
def CalConvoluationLayer(img_input,Convlayer_dict):
    '''
    计算单个图层的计算结果
    args:
        img_input:tensor 1,c,h,w 处理图像
        Convlayer_dict:moudle 加速
    convlayers['conv1']={'conv':conv_modules['conv_1'],
                        'batchnorm':conv_modules['batchNorm2d_1'],
                        'ReLU':conv_modules['ReLU_1']}
    '''
    result={'img_input':img_input.detach().cpu().numpy()}
    # 计算卷积层
    result['conv_result']=CalulationConvolution(img_input,Convlayer_dict['conv'].weight)
    result['batchnorm_result']=BatchNorm(torch.from_numpy( result['conv_result']['output']).cuda().float(),Convlayer_dict['batchnorm'])
    result['ReLu_result']=ReLuCalucation(torch.from_numpy(result['batchnorm_result']['output']),Convlayer_dict['ReLU'])
    return result

def CaluationConvlayers(img_input,convlayer_dict):
    '''
    计算前两层结果
    '''
    c,h,w=img_input.shape
    img_input=img_input.reshape(1,c,h,w)
    img_input=torch.from_numpy(img_input)
    conv1_result=CalConvoluationLayer(img_input,convlayer_dict['conv1'])
    conv2_result=CalConvoluationLayer(torch.from_numpy(conv1_result['ReLu_result'])  ,convlayer_dict['conv2'])
    calucation_conv_result={'conv1':conv1_result,'conv2':conv2_result} # 获取计算的结果
    return calucation_conv_result

def SaveCaluationResult(root_dir,cal_result):
    # 构建保存路径
    CreateFolder(root_dir)
    for layername in cal_result:
        layername_path=os.path.join(root_dir,layername)
        CreateFolder(layername_path)
        # 保存输入文件
        np.save(os.path.join(layername_path,'input_img.npy'),cal_result[layername]['img_input'])
        # 存储卷积层计算结果
        conv_path=os.path.join(layername_path,'conv')
        CreateFolder(conv_path)
        conv_input_img=cal_result[layername]['conv_result']['img']
        np.save(os.path.join(conv_path,'input_img.npy'),conv_input_img)
        # 存储卷积层的核权重
        np.save(os.path.join(conv_path,'kernel_weight.npy'),cal_result[layername]['conv_result']['kernels'])
        # 处理单个通道
        conv_kernel_sig=os.path.join(conv_path,'kernel_sig')
        CreateFolder(conv_kernel_sig)
        for kernel_index in cal_result[layername]['conv_result']['kernel_sig']:
            kernel_sig_temp=cal_result[layername]['conv_result']['kernel_sig'][kernel_index]
            conv_kernel_sig_index_path=os.path.join(conv_kernel_sig,str(kernel_index))
            CreateFolder(conv_kernel_sig_index_path)
            np.save(os.path.join(conv_kernel_sig_index_path,'kernel_weight.npy'),kernel_sig_temp['kernel'])
            # 记录当前各个特征的卷积核
            for channel_index in kernel_sig_temp['kernel_result']:
                channel_path=os.path.join(conv_kernel_sig_index_path,'channel_{}'.format(channel_index))
                CreateFolder(channel_path)
                np.save(os.path.join(channel_path,'kernel_weight.npy'),kernel_sig_temp['kernel_result'][channel_index]['weight'])
                np.save(os.path.join(channel_path,'output.npy'),kernel_sig_temp['kernel_result'][channel_index]['result'])
                pass
            # 记录各个波段贡献度计算结果
            np.save(os.path.join(conv_kernel_sig_index_path,'contribution.npy'),kernel_sig_temp['contribution'])
        # 保存最终的输出
        np.save(os.path.join(conv_path,'output.npy'),cal_result[layername]['conv_result']['output'])
        pass
    pass